import os
from pathlib import Path

import sklearn.datasets as data
import torch
from sklearn.model_selection import train_test_split

from . import TorchTest, TorchTrain
from .TorchRegressionModule_iris import RegressionModule


def run():
    iris_data = data.load_iris()
    x = iris_data['data']
    y = iris_data['target']
    x_train, x_test, y_train, y_test = train_test_split(x, y, test_size=0.33, random_state=None)

    save_path = os.path.abspath('') + '/abc'
    if Path(save_path).is_file():
        model = torch.load(save_path)
    else:
        model = RegressionModule()
        TorchTrain.train(model, torch.tensor(x_train, dtype=torch.float), torch.tensor(y_train, dtype=torch.long),
                         save_path)

    for parameter in model.parameters():
        print(parameter)

    model.eval()

    TorchTest.test(model,
                   torch.tensor(x_test, dtype=torch.float),
                   torch.tensor(y_test, dtype=torch.long))
